#!/usr/bin/env python3
import argparse
import os
import sys
from time import sleep

import grpc

sys.path.append(
    os.path.join(os.path.dirname(os.path.abspath(__file__)),
                 '../../utils/'))
import p4runtime_lib.bmv2
import p4runtime_lib.helper
from p4runtime_lib.switch import ShutdownAllSwitchConnections

def writeIpv4Rules(p4info_helper, switch, host_ips, host_macs, switch_ports):
    """
    为每个交换机安装IPv4转发规则
    
    :param p4info_helper: P4Info helper
    :param switch: 交换机连接对象
    :param host_ips: 主机IP地址列表
    :param host_macs: 主机MAC地址列表  
    :param switch_ports: 主机连接的端口列表
    """
    for i, (host_ip, host_mac, port) in enumerate(zip(host_ips, host_macs, switch_ports)):
        table_entry = p4info_helper.buildTableEntry(
            table_name="MyIngress.ipv4_lpm",
            match_fields={
                "hdr.ipv4.dstAddr": (host_ip, 32)
            },
            action_name="MyIngress.ipv4_forward",
            action_params={
                "dstAddr": host_mac,
                "port": port
            })
        switch.WriteTableEntry(table_entry)
        print("Installed IPv4 forwarding rule on %s for %s -> port %d" % 
              (switch.name, host_ip, port))

def writeSrcRoutingRules(p4info_helper, switch, route_entries):
    """
    安装源路由规则（如果需要的话）
    在实际场景中，源路由路径通常由控制器动态计算
    """
    # 这里可以添加源路由相关的表项
    # 由于源路由逻辑主要在数据平面处理，控制平面主要负责路径计算
    pass

def readTableRules(p4info_helper, sw):
    """
    读取交换机上的所有表规则
    """
    print('\n----- Reading tables rules for %s -----' % sw.name)
    for response in sw.ReadTableEntries():
        for entity in response.entities:
            entry = entity.table_entry
            table_name = p4info_helper.get_tables_name(entry.table_id)
            print('%s: ' % table_name, end=' ')
            for m in entry.match:
                print(p4info_helper.get_match_field_name(table_name, m.field_id), end=' ')
                print('%r' % (p4info_helper.get_match_field_value(m),), end=' ')
            action = entry.action.action
            action_name = p4info_helper.get_actions_name(action.action_id)
            print('->', action_name, end=' ')
            for p in action.params:
                print(p4info_helper.get_action_param_name(action_name, p.param_id), end=' ')
                print('%r' % p.value, end=' ')
            print()

def main(p4info_file_path, bmv2_file_path):
    # 实例化P4Runtime helper
    p4info_helper = p4runtime_lib.helper.P4InfoHelper(p4info_file_path)

    try:
        # 创建交换机连接
        s1 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s1',
            address='127.0.0.1:50051',
            device_id=0,
            proto_dump_file='logs/s1-p4runtime-requests.txt')
        s2 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s2',
            address='127.0.0.1:50052',
            device_id=1,
            proto_dump_file='logs/s2-p4runtime-requests.txt')

        # 发送master仲裁更新
        s1.MasterArbitrationUpdate()
        s2.MasterArbitrationUpdate()

        # 安装P4程序
        s1.SetForwardingPipelineConfig(p4info=p4info_helper.p4info,
                                       bmv2_json_file_path=bmv2_file_path)
        print("Installed P4 Program using SetForwardingPipelineConfig on s1")
        s2.SetForwardingPipelineConfig(p4info=p4info_helper.p4info,
                                       bmv2_json_file_path=bmv2_file_path)
        print("Installed P4 Program using SetForwardingPipelineConfig on s2")

        # 为s1安装IPv4转发规则
        # 假设拓扑：h1(10.0.1.1) -- s1 -- s2 -- h2(10.0.2.2)
        writeIpv4Rules(p4info_helper, s1, 
                      host_ips=["10.0.1.1", "10.0.2.2"], 
                      host_macs=["08:00:00:00:01:11", "08:00:00:00:02:22"],
                      switch_ports=[1, 2])  # 端口1连接h1，端口2连接s2

        # 为s2安装IPv4转发规则
        writeIpv4Rules(p4info_helper, s2,
                      host_ips=["10.0.1.1", "10.0.2.2"],
                      host_macs=["08:00:00:00:01:11", "08:00:00:00:02:22"], 
                      switch_ports=[2, 1])  # 端口2连接s1，端口1连接h2

        # 读取表规则
        readTableRules(p4info_helper, s1)
        readTableRules(p4info_helper, s2)

        print("\nSetup complete! Switch now supports both IPv4 forwarding and source routing.")
        print("Use send.py to send IPv4 packets or source-routed packets.")

        # 保持运行
        while True:
            sleep(10)

    except KeyboardInterrupt:
        print(" Shutting down.")
    except grpc.RpcError as e:
        print("gRPC Error:", e.details())

    ShutdownAllSwitchConnections()

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='P4Runtime Controller')
    parser.add_argument('--p4info', help='p4info proto in text format from p4c',
                        type=str, action="store", required=False,
                        default='./build/source_routing.p4.p4info.txt')
    parser.add_argument('--bmv2-json', help='BMv2 JSON file from p4c',
                        type=str, action="store", required=False,
                        default='./build/source_routing.json')
    args = parser.parse_args()

    if not os.path.exists(args.p4info):
        parser.print_help()
        print("\np4info file not found: %s\nHave you run 'make'?" % args.p4info)
        parser.exit(1)
    if not os.path.exists(args.bmv2_json):
        parser.print_help()
        print("\nBMv2 JSON file not found: %s\nHave you run 'make'?" % args.bmv2_json)
        parser.exit(1)
    main(args.p4info, args.bmv2_json)